{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vXLA5InzXydn"
      },
      "source": [
        "##### Copyright 2019 The TensorFlow Authors."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "cellView": "form",
        "id": "RuRlpLL-X0R_"
      },
      "outputs": [],
      "source": [
        "#@title Licensed under the Apache License, Version 2.0 (the \"License\");\n",
        "# you may not use this file except in compliance with the License.\n",
        "# You may obtain a copy of the License at\n",
        "#\n",
        "# https://www.apache.org/licenses/LICENSE-2.0\n",
        "#\n",
        "# Unless required by applicable law or agreed to in writing, software\n",
        "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
        "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
        "# See the License for the specific language governing permissions and\n",
        "# limitations under the License."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "1mLJmVotXs64"
      },
      "source": [
        "# Fine-tuning a BERT model"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "hYEwGTeCXnnX"
      },
      "source": [
        "\u003ctable class=\"tfo-notebook-buttons\" align=\"left\"\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca target=\"_blank\" href=\"https://www.tensorflow.org/tfmodels/nlp/fine_tune_bert\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/tf_logo_32px.png\" /\u003eView on TensorFlow.org\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca target=\"_blank\" href=\"https://colab.research.google.com/github/tensorflow/models/blob/master/docs/nlp/fine_tune_bert.ipynb\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/colab_logo_32px.png\" /\u003eRun in Google Colab\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca target=\"_blank\" href=\"https://github.com/tensorflow/models/blob/master/docs/nlp/fine_tune_bert.ipynb\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/GitHub-Mark-32px.png\" /\u003eView source on GitHub\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca href=\"https://storage.googleapis.com/tensorflow_docs/models/docs/nlp/fine_tune_bert.ipynb\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/download_logo_32px.png\" /\u003eDownload notebook\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "  \u003ctd\u003e\n",
        "    \u003ca href=\"https://tfhub.dev/google/collections/bert\"\u003e\u003cimg src=\"https://www.tensorflow.org/images/hub_logo_32px.png\" /\u003eSee TF Hub model\u003c/a\u003e\n",
        "  \u003c/td\u003e\n",
        "\u003c/table\u003e"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "YN2ACivEPxgD"
      },
      "source": [
        "This tutorial demonstrates how to fine-tune a [Bidirectional Encoder Representations from Transformers (BERT)](https://arxiv.org/abs/1810.04805) (Devlin et al., 2018) model using [TensorFlow Model Garden](https://github.com/tensorflow/models).\n",
        "\n",
        "You can also find the pre-trained BERT model used in this tutorial on [TensorFlow Hub (TF Hub)](https://tensorflow.org/hub). For concrete examples of how to use the models from TF Hub, refer to the [Solve Glue tasks using BERT](https://www.tensorflow.org/text/tutorials/bert_glue) tutorial. If you're just trying to fine-tune a model, the TF Hub tutorial is a good starting point.\n",
        "\n",
        "On the other hand, if you're interested in deeper customization, follow this tutorial. It shows how to do a lot of things manually, so you can learn how you can customize the workflow from data preprocessing to training, exporting and saving the model."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "s2d9S2CSSO1z"
      },
      "source": [
        "## Setup"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "69de3375e32a"
      },
      "source": [
        "### Install pip packages"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fsACVQpVSifi"
      },
      "source": [
        "Start by installing the TensorFlow Text and Model Garden pip packages.\n",
        "\n",
        "*  `tf-models-official` is the TensorFlow Model Garden package. Note that it may not include the latest changes in the `tensorflow_models` GitHub repo. To include the latest changes, you may install `tf-models-nightly`, which is the nightly Model Garden package created daily automatically.\n",
        "*  pip will install all models and dependencies automatically."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "sE6XUxLOf1s-"
      },
      "outputs": [],
      "source": [
        "!pip install -q opencv-python"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "yic2y7_o-BCC"
      },
      "outputs": [],
      "source": [
        "!pip install -q -U \"tensorflow-text==2.11.*\""
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "NvNr2svBM-p3"
      },
      "outputs": [],
      "source": [
        "!pip install -q tf-models-official"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "U-7qPCjWUAyy"
      },
      "source": [
        "### Import libraries"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "lXsXev5MNr20"
      },
      "outputs": [],
      "source": [
        "import os\n",
        "\n",
        "import numpy as np\n",
        "import matplotlib.pyplot as plt\n",
        "\n",
        "import tensorflow as tf\n",
        "import tensorflow_models as tfm\n",
        "import tensorflow_hub as hub\n",
        "import tensorflow_datasets as tfds\n",
        "tfds.disable_progress_bar()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mbanlzTvJBsz"
      },
      "source": [
        "### Resources"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "PpW0x8TpR8DT"
      },
      "source": [
        "The following directory contains the BERT model's configuration, vocabulary, and a pre-trained checkpoint used in this tutorial:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "vzRHOLciR8eq"
      },
      "outputs": [],
      "source": [
        "gs_folder_bert = \"gs://cloud-tpu-checkpoints/bert/v3/uncased_L-12_H-768_A-12\"\n",
        "tf.io.gfile.listdir(gs_folder_bert)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Qv6abtRvH4xO"
      },
      "source": [
        "## Load and preprocess the dataset\n",
        "\n",
        "This example uses the GLUE (General Language Understanding Evaluation) MRPC (Microsoft Research Paraphrase Corpus) [dataset from TensorFlow Datasets (TFDS)](https://www.tensorflow.org/datasets/catalog/glue#gluemrpc).\n",
        "\n",
        "This dataset is not set up such that it can be directly fed into the BERT model. The following section handles the necessary preprocessing."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "28DvUhC1YUiB"
      },
      "source": [
        "### Get the dataset from TensorFlow Datasets\n",
        "\n",
        "The GLUE MRPC (Dolan and Brockett, 2005) dataset is a corpus of sentence pairs automatically extracted from online news sources, with human annotations for whether the sentences in the pair are semantically equivalent. It has the following attributes:\n",
        "\n",
        "*   Number of labels: 2\n",
        "*   Size of training dataset: 3668\n",
        "*   Size of evaluation dataset: 408\n",
        "*   Maximum sequence length of training and evaluation dataset: 128\n",
        "\n",
        "Begin by loading the MRPC dataset from TFDS:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Ijikx5OsH9AT"
      },
      "outputs": [],
      "source": [
        "batch_size=32\n",
        "glue, info = tfds.load('glue/mrpc',\n",
        "                       with_info=True,\n",
        "                       batch_size=32)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "QcMTJU4N7VX-"
      },
      "outputs": [],
      "source": [
        "glue"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ZgBg2r2nYT-K"
      },
      "source": [
        "The `info` object describes the dataset and its features:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "IQrHxv7W7jH5"
      },
      "outputs": [],
      "source": [
        "info.features"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "vhsVWYNxazz5"
      },
      "source": [
        "The two classes are:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "n0gfc_VTayfQ"
      },
      "outputs": [],
      "source": [
        "info.features['label'].names"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "38zJcap6xkbC"
      },
      "source": [
        "Here is one example from the training set:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "xON_i6SkwApW"
      },
      "outputs": [],
      "source": [
        "example_batch = next(iter(glue['train']))\n",
        "\n",
        "for key, value in example_batch.items():\n",
        "  print(f\"{key:9s}: {value[0].numpy()}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "R9vEWgKA4SxV"
      },
      "source": [
        "### Preprocess the data\n",
        "\n",
        "The keys `\"sentence1\"` and `\"sentence2\"` in the GLUE MRPC dataset contain two input sentences for each example.\n",
        "\n",
        "Because the BERT model from the Model Garden doesn't take raw text as input, two things need to happen first:\n",
        "\n",
        "1. The text needs to be _tokenized_ (split into word pieces) and converted to _indices_.\n",
        "2. Then, the _indices_ need to be packed into the format that the model expects."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "9fbTyfJpNr7x"
      },
      "source": [
        "#### The BERT tokenizer"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wqeN54S61ZKQ"
      },
      "source": [
        "To fine tune a pre-trained language model from the Model Garden, such as BERT, you need to make sure that you're using exactly the same tokenization, vocabulary, and index mapping as used during training.\n",
        "\n",
        "The following code rebuilds the tokenizer that was used by the base model using the Model Garden's `tfm.nlp.layers.FastWordpieceBertTokenizer` layer:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "-DK4q5wEBmlB"
      },
      "outputs": [],
      "source": [
        "tokenizer = tfm.nlp.layers.FastWordpieceBertTokenizer(\n",
        "    vocab_file=os.path.join(gs_folder_bert, \"vocab.txt\"),\n",
        "    lower_case=True)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "zYHDSquU2lDU"
      },
      "source": [
        "Let's tokenize a test sentence:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "L_OfOYPg853R"
      },
      "outputs": [],
      "source": [
        "tokens = tokenizer(tf.constant([\"Hello TensorFlow!\"]))\n",
        "tokens"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MfjaaMYy5Gt8"
      },
      "source": [
        "Learn more about the tokenization process in the [Subword tokenization](https://www.tensorflow.org/text/guide/subwords_tokenizer) and [Tokenizing with TensorFlow Text](https://www.tensorflow.org/text/guide/tokenizers) guides."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "wd1b09OO5GJl"
      },
      "source": [
        "#### Pack the inputs"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "62UTWLQd9-LB"
      },
      "source": [
        "TensorFlow Model Garden's BERT model doesn't just take the tokenized strings as input. It also expects these to be packed into a particular format. `tfm.nlp.layers.BertPackInputs` layer can handle the conversion from _a list of tokenized sentences_ to the input format expected by the Model Garden's BERT model.\n",
        "\n",
        "`tfm.nlp.layers.BertPackInputs` packs the two input sentences (per example in the MRCP dataset) concatenated together. This input is expected to start with a `[CLS]` \"This is a classification problem\" token, and each sentence should end with a `[SEP]` \"Separator\" token.\n",
        "\n",
        "Therefore, the `tfm.nlp.layers.BertPackInputs` layer's constructor takes the `tokenizer`'s special tokens as an argument. It also needs to know the indices of the tokenizer's special tokens."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5iroDlrFDRcF"
      },
      "outputs": [],
      "source": [
        "special = tokenizer.get_special_tokens_dict()\n",
        "special"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "b71HarkuG92H"
      },
      "outputs": [],
      "source": [
        "max_seq_length = 128\n",
        "\n",
        "packer = tfm.nlp.layers.BertPackInputs(\n",
        "    seq_length=max_seq_length,\n",
        "    special_tokens_dict = tokenizer.get_special_tokens_dict())"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CZlSZbYd6liN"
      },
      "source": [
        "The `packer` takes a list of tokenized sentences as input. For example:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "27dU_VkJHc9S"
      },
      "outputs": [],
      "source": [
        "sentences1 = [\"hello tensorflow\"]\n",
        "tok1 = tokenizer(sentences1)\n",
        "tok1"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "LURHmNOSHnWN"
      },
      "outputs": [],
      "source": [
        "sentences2 = [\"goodbye tensorflow\"]\n",
        "tok2 = tokenizer(sentences2)\n",
        "tok2"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "r8bvB8gI8BqP"
      },
      "source": [
        "Then, it returns a dictionary containing three outputs:\n",
        "\n",
        "- `input_word_ids`: The tokenized sentences packed together.\n",
        "- `input_mask`: The mask indicating which locations are valid in the other outputs.\n",
        "- `input_type_ids`: Indicating which sentence each token belongs to."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "YsIDTOMJHrUQ"
      },
      "outputs": [],
      "source": [
        "packed = packer([tok1, tok2])\n",
        "\n",
        "for key, tensor in packed.items():\n",
        "  print(f\"{key:15s}: {tensor[:, :12]}\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "red4tRcq74Qc"
      },
      "source": [
        "#### Put it all together\n",
        "\n",
        "Combine these two parts into a `keras.layers.Layer` that can be attached to your model:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9Qtz-tv-6nz6"
      },
      "outputs": [],
      "source": [
        "class BertInputProcessor(tf.keras.layers.Layer):\n",
        "  def __init__(self, tokenizer, packer):\n",
        "    super().__init__()\n",
        "    self.tokenizer = tokenizer\n",
        "    self.packer = packer\n",
        "\n",
        "  def call(self, inputs):\n",
        "    tok1 = self.tokenizer(inputs['sentence1'])\n",
        "    tok2 = self.tokenizer(inputs['sentence2'])\n",
        "\n",
        "    packed = self.packer([tok1, tok2])\n",
        "\n",
        "    if 'label' in inputs:\n",
        "      return packed, inputs['label']\n",
        "    else:\n",
        "      return packed"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rdy9wp499btU"
      },
      "source": [
        "But for now just apply it to the dataset using `Dataset.map`, since the dataset you loaded from TFDS is a `tf.data.Dataset` object:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "qmyh76AL7VAs"
      },
      "outputs": [],
      "source": [
        "bert_inputs_processor = BertInputProcessor(tokenizer, packer)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "B8SSCtDe9MCk"
      },
      "outputs": [],
      "source": [
        "glue_train = glue['train'].map(bert_inputs_processor).prefetch(1)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "KXpiDosO9rkY"
      },
      "source": [
        "Here is an example batch from the processed dataset:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "ffNvDE6t9rP-"
      },
      "outputs": [],
      "source": [
        "example_inputs, example_labels = next(iter(glue_train))"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "5sxtTuUi-bXt"
      },
      "outputs": [],
      "source": [
        "example_inputs"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "wP4z_-9a-dFk"
      },
      "outputs": [],
      "source": [
        "example_labels"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "jyjTdGpFhO_1"
      },
      "outputs": [],
      "source": [
        "for key, value in example_inputs.items():\n",
        "  print(f'{key:15s} shape: {value.shape}')\n",
        "\n",
        "print(f'{\"labels\":15s} shape: {example_labels.shape}')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mkGHN_FK-50U"
      },
      "source": [
        "The `input_word_ids` contain the token IDs:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "eGL1_ktWLcgF"
      },
      "outputs": [],
      "source": [
        "plt.pcolormesh(example_inputs['input_word_ids'])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ulNZ4U96-8JZ"
      },
      "source": [
        "The mask allows the model to cleanly differentiate between the content and the padding. The mask has the same shape as the `input_word_ids`, and contains a `1` anywhere the `input_word_ids` is not padding."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "zB7mW7DGK3rW"
      },
      "outputs": [],
      "source": [
        "plt.pcolormesh(example_inputs['input_mask'])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "rxLenwAvCkBf"
      },
      "source": [
        "The \"input type\" also has the same shape, but inside the non-padded region, contains a `0` or a `1` indicating which sentence the token is a part of."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "2CetH_5C9P2m"
      },
      "outputs": [],
      "source": [
        "plt.pcolormesh(example_inputs['input_type_ids'])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "pxHHeyei_sb9"
      },
      "source": [
        "Apply the same preprocessing to the validation and test subsets of the GLUE MRPC dataset:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "yuLKxf6zHxw-"
      },
      "outputs": [],
      "source": [
        "glue_validation = glue['validation'].map(bert_inputs_processor).prefetch(1)\n",
        "glue_test = glue['test'].map(bert_inputs_processor).prefetch(1)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "FSwymsbkbLDA"
      },
      "source": [
        "## Build, train and export the model"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bxxO3pJCEM9p"
      },
      "source": [
        "Now that you have formatted the data as expected, you can start working on building and training the model."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Efrj3Cn1kLAp"
      },
      "source": [
        "### Build the model\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xxpOY5r2Ayq6"
      },
      "source": [
        "The first step is to download the configuration file—`config_dict`—for the pre-trained BERT model:\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "v7ap0BONSJuz"
      },
      "outputs": [],
      "source": [
        "import json\n",
        "\n",
        "bert_config_file = os.path.join(gs_folder_bert, \"bert_config.json\")\n",
        "config_dict = json.loads(tf.io.gfile.GFile(bert_config_file).read())\n",
        "config_dict"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "pKaEaKJSX85J"
      },
      "outputs": [],
      "source": [
        "encoder_config = tfm.nlp.encoders.EncoderConfig({\n",
        "    'type':'bert',\n",
        "    'bert': config_dict\n",
        "})"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "LbgzWukNSqOS"
      },
      "outputs": [],
      "source": [
        "bert_encoder = tfm.nlp.encoders.build_encoder(encoder_config)\n",
        "bert_encoder"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "96ldxDSwkVkj"
      },
      "source": [
        "The configuration file defines the core BERT model from the Model Garden, which is a Keras model that predicts the outputs of `num_classes` from the inputs with maximum sequence length `max_seq_length`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "cH682__U0FBv"
      },
      "outputs": [],
      "source": [
        "bert_classifier = tfm.nlp.models.BertClassifier(network=bert_encoder, num_classes=2)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "sFmVG4SKZAw8"
      },
      "source": [
        "Run it on a test batch of data 10 examples from the training set. The output is the logits for the two classes:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "VTjgPbp4ZDKo"
      },
      "outputs": [],
      "source": [
        "bert_classifier(\n",
        "    example_inputs, training=True).numpy()[:10]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Q0NTdwZsQK8n"
      },
      "source": [
        "The `TransformerEncoder` in the center of the classifier above **is** the `bert_encoder`.\n",
        "\n",
        "If you inspect the encoder, notice the stack of `Transformer` layers connected to those same three inputs:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "8L__-erBwLIQ"
      },
      "outputs": [],
      "source": [
        "tf.keras.utils.plot_model(bert_encoder, show_shapes=True, dpi=48)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "mKAvkQc3heSy"
      },
      "source": [
        "### Restore the encoder weights\n",
        "\n",
        "When built, the encoder is randomly initialized. Restore the encoder's weights from the checkpoint:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "97Ll2Gichd_Y"
      },
      "outputs": [],
      "source": [
        "checkpoint = tf.train.Checkpoint(encoder=bert_encoder)\n",
        "checkpoint.read(\n",
        "    os.path.join(gs_folder_bert, 'bert_model.ckpt')).assert_consumed()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2oHOql35k3Dd"
      },
      "source": [
        "Note: The pretrained `TransformerEncoder` is also available on [TensorFlow Hub](https://tensorflow.org/hub). Go to the [TF Hub appendix](#hub_bert) for details."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "115caFLMk-_l"
      },
      "source": [
        "### Set up the optimizer\n",
        "\n",
        "BERT typically uses the Adam optimizer with weight decay—[AdamW](https://arxiv.org/abs/1711.05101) (`tf.keras.optimizers.experimental.AdamW`).\n",
        "It also employs a learning rate schedule that first warms up from 0 and then decays to 0:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "c0jBycPDtkxR"
      },
      "outputs": [],
      "source": [
        "# Set up epochs and steps\n",
        "epochs = 5\n",
        "batch_size = 32\n",
        "eval_batch_size = 32\n",
        "\n",
        "train_data_size = info.splits['train'].num_examples\n",
        "steps_per_epoch = int(train_data_size / batch_size)\n",
        "num_train_steps = steps_per_epoch * epochs\n",
        "warmup_steps = int(0.1 * num_train_steps)\n",
        "initial_learning_rate=2e-5"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "GFankgHK0Rvh"
      },
      "source": [
        "Linear decay from `initial_learning_rate` to zero over `num_train_steps`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "qWSyT8P2j4mV"
      },
      "outputs": [],
      "source": [
        "linear_decay = tf.keras.optimizers.schedules.PolynomialDecay(\n",
        "    initial_learning_rate=initial_learning_rate,\n",
        "    end_learning_rate=0,\n",
        "    decay_steps=num_train_steps)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "anZPZPAP0Y3n"
      },
      "source": [
        "Warmup to that value over `warmup_steps`:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "z_AsVCiRkoN1"
      },
      "outputs": [],
      "source": [
        "warmup_schedule = tfm.optimization.lr_schedule.LinearWarmup(\n",
        "    warmup_learning_rate = 0,\n",
        "    after_warmup_lr_sched = linear_decay,\n",
        "    warmup_steps = warmup_steps\n",
        ")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "arfbaK6t0kH_"
      },
      "source": [
        "The overall schedule looks like this:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "rYZGunhqbGUZ"
      },
      "outputs": [],
      "source": [
        "x = tf.linspace(0, num_train_steps, 1001)\n",
        "y = [warmup_schedule(xi) for xi in x]\n",
        "plt.plot(x,y)\n",
        "plt.xlabel('Train step')\n",
        "plt.ylabel('Learning rate')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "bjsmG_fm0opn"
      },
      "source": [
        "Use `tf.keras.optimizers.experimental.AdamW` to instantiate the optimizer with that schedule:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "R8pTNuKIw1dA"
      },
      "outputs": [],
      "source": [
        "optimizer = tf.keras.optimizers.experimental.Adam(\n",
        "    learning_rate = warmup_schedule)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "78FEUOOEkoP0"
      },
      "source": [
        "### Train the model"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "OTNcA0O0nSq9"
      },
      "source": [
        "Set the metric as accuracy and the loss as sparse categorical cross-entropy. Then, compile and train the BERT classifier:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "d5FeL0b6j7ky"
      },
      "outputs": [],
      "source": [
        "metrics = [tf.keras.metrics.SparseCategoricalAccuracy('accuracy', dtype=tf.float32)]\n",
        "loss = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)\n",
        "\n",
        "bert_classifier.compile(\n",
        "    optimizer=optimizer,\n",
        "    loss=loss,\n",
        "    metrics=metrics)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "CsrylctIj_Xy"
      },
      "outputs": [],
      "source": [
        "bert_classifier.evaluate(glue_validation)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "hgPPc2oNmcVZ"
      },
      "outputs": [],
      "source": [
        "bert_classifier.fit(\n",
        "      glue_train,\n",
        "      validation_data=(glue_validation),\n",
        "      batch_size=32,\n",
        "      epochs=epochs)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IFtKFWbNKb0u"
      },
      "source": [
        "Now run the fine-tuned model on a custom example to see that it works.\n",
        "\n",
        "Start by encoding some sentence pairs:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "S1sdW6lLWaEi"
      },
      "outputs": [],
      "source": [
        "my_examples = {\n",
        "        'sentence1':[\n",
        "            'The rain in Spain falls mainly on the plain.',\n",
        "            'Look I fine tuned BERT.'],\n",
        "        'sentence2':[\n",
        "            'It mostly rains on the flat lands of Spain.',\n",
        "            'Is it working? This does not match.']\n",
        "    }"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "7ynJibkBRTJF"
      },
      "source": [
        "The model should report class `1` \"match\" for the first example and class `0` \"no-match\" for the second:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "umo0ttrgRYIM"
      },
      "outputs": [],
      "source": [
        "ex_packed = bert_inputs_processor(my_examples)\n",
        "my_logits = bert_classifier(ex_packed, training=False)\n",
        "\n",
        "result_cls_ids = tf.argmax(my_logits)\n",
        "result_cls_ids"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "HNdmOEHKT7e8"
      },
      "outputs": [],
      "source": [
        "tf.gather(tf.constant(info.features['label'].names), result_cls_ids)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "fVo_AnT0l26j"
      },
      "source": [
        "### Export the model\n",
        "\n",
        "Often the goal of training a model is to _use_ it for something outside of the Python process that created it. You can do this by exporting the model using `tf.saved_model`. (Learn more in the [Using the SavedModel format](https://www.tensorflow.org/guide/saved_model) guide and the [Save and load a model using a distribution strategy](https://www.tensorflow.org/tutorials/distribute/save_and_load) tutorial.)\n",
        "\n",
        "First, build a wrapper class to export the model. This wrapper does two things:\n",
        "\n",
        "- First it packages `bert_inputs_processor` and `bert_classifier` together into a single `tf.Module`, so you can export all the functionalities.\n",
        "- Second it defines a `tf.function` that implements the end-to-end execution of the model.\n",
        "\n",
        "Setting the `input_signature` argument of `tf.function` lets you define a fixed signature for the `tf.function`. This can be less surprising than the default automatic retracing behavior."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "78h83mlt9wpY"
      },
      "outputs": [],
      "source": [
        "class ExportModel(tf.Module):\n",
        "  def __init__(self, input_processor, classifier):\n",
        "    self.input_processor = input_processor\n",
        "    self.classifier = classifier\n",
        "\n",
        "  @tf.function(input_signature=[{\n",
        "      'sentence1': tf.TensorSpec(shape=[None], dtype=tf.string),\n",
        "      'sentence2': tf.TensorSpec(shape=[None], dtype=tf.string)}])\n",
        "  def __call__(self, inputs):\n",
        "    packed = self.input_processor(inputs)\n",
        "    logits =  self.classifier(packed, training=False)\n",
        "    result_cls_ids = tf.argmax(logits)\n",
        "    return {\n",
        "        'logits': logits,\n",
        "        'class_id': result_cls_ids,\n",
        "        'class': tf.gather(\n",
        "            tf.constant(info.features['label'].names),\n",
        "            result_cls_ids)\n",
        "    }"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "qnxysGUfIgFQ"
      },
      "source": [
        "Create an instance of this export-model and save it:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "TmHW9DEFUZ0X"
      },
      "outputs": [],
      "source": [
        "export_model = ExportModel(bert_inputs_processor, bert_classifier)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Nl5x6nElZqkP"
      },
      "outputs": [],
      "source": [
        "import tempfile\n",
        "export_dir=tempfile.mkdtemp(suffix='_saved_model')\n",
        "tf.saved_model.save(export_model, export_dir=export_dir,\n",
        "                    signatures={'serving_default': export_model.__call__})"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Pd8B5dy-ImDJ"
      },
      "source": [
        "Reload the model and compare the results to the original:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9cAhHySVXHD5"
      },
      "outputs": [],
      "source": [
        "original_logits = export_model(my_examples)['logits']"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "H9cAcYwfW2fy"
      },
      "outputs": [],
      "source": [
        "reloaded = tf.saved_model.load(export_dir)\n",
        "reloaded_logits = reloaded(my_examples)['logits']"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "y_ACvKPsVUXC"
      },
      "outputs": [],
      "source": [
        "# The results are identical:\n",
        "print(original_logits.numpy())\n",
        "print()\n",
        "print(reloaded_logits.numpy())"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "lBlPP20dXPFR"
      },
      "outputs": [],
      "source": [
        "print(np.mean(abs(original_logits - reloaded_logits)))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "CPsg7dZwfBM2"
      },
      "source": [
        "Congratulations! You've used `tensorflow_models` to build a BERT-classifier, train it, and export for later use."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "eQceYqRFT_Eg"
      },
      "source": [
        "## Optional: BERT on TF Hub"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QbklKt-w_CiI"
      },
      "source": [
        "\u003ca id=\"hub_bert\"\u003e\u003c/a\u003e\n",
        "\n",
        "\n",
        "You can get the BERT model off the shelf from [TF Hub](https://tfhub.dev/). There are [many versions available along with their input preprocessors](https://tfhub.dev/google/collections/bert/1).\n",
        "\n",
        "This example uses [a small version of BERT from TF Hub](https://tfhub.dev/tensorflow/small_bert/bert_en_uncased_L-2_H-128_A-2/2) that was pre-trained using the English Wikipedia and BooksCorpus datasets, similar to the [original implementation](https://arxiv.org/abs/1908.08962) (Turc et al., 2019).\n",
        "\n",
        "Start by importing TF Hub:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "GDWrHm0BGpbX"
      },
      "outputs": [],
      "source": [
        "import tensorflow_hub as hub"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "f02f38f83ac4"
      },
      "source": [
        "Select the input preprocessor and the model from TF Hub and wrap them as `hub.KerasLayer` layers:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "lo6479At4sP1"
      },
      "outputs": [],
      "source": [
        "# Always make sure you use the right preprocessor.\n",
        "hub_preprocessor = hub.KerasLayer(\n",
        "    \"https://tfhub.dev/tensorflow/bert_en_uncased_preprocess/3\")\n",
        "\n",
        "# This is a really small BERT.\n",
        "hub_encoder = hub.KerasLayer(f\"https://tfhub.dev/tensorflow/small_bert/bert_en_uncased_L-2_H-128_A-2/2\",\n",
        "                             trainable=True)\n",
        "\n",
        "print(f\"The Hub encoder has {len(hub_encoder.trainable_variables)} trainable variables\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "iTzF574wivQv"
      },
      "source": [
        "Test run the preprocessor on a batch of data:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "GOASSKR5R3-N"
      },
      "outputs": [],
      "source": [
        "hub_inputs = hub_preprocessor(['Hello TensorFlow!'])\n",
        "{key: value[0, :10].numpy() for key, value in hub_inputs.items()} "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "XEcYrCR45Uwo"
      },
      "outputs": [],
      "source": [
        "result = hub_encoder(\n",
        "    inputs=hub_inputs,\n",
        "    training=False,\n",
        ")\n",
        "\n",
        "print(\"Pooled output shape:\", result['pooled_output'].shape)\n",
        "print(\"Sequence output shape:\", result['sequence_output'].shape)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "cjojn8SmLSRI"
      },
      "source": [
        "At this point it would be simple to add a classification head yourself.\n",
        "\n",
        "The Model Garden `tfm.nlp.models.BertClassifier` class can also build a classifier onto the TF Hub encoder:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "9nTDaApyLR70"
      },
      "outputs": [],
      "source": [
        "hub_classifier = tfm.nlp.models.BertClassifier(\n",
        "    bert_encoder,\n",
        "    num_classes=2,\n",
        "    dropout_rate=0.1,\n",
        "    initializer=tf.keras.initializers.TruncatedNormal(\n",
        "        stddev=0.02))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xMJX3wV0_v7I"
      },
      "source": [
        "The one downside to loading this model from TF Hub is that the structure of internal Keras layers is not restored. This makes it more difficult to inspect or modify the model.\n",
        "\n",
        "The BERT encoder model—`hub_classifier`—is now a single layer."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "u_IqwXjRV1vd"
      },
      "source": [
        "For concrete examples of this approach, refer to [Solve Glue tasks using BERT](https://www.tensorflow.org/text/tutorials/bert_glue)."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ji3tdLz101km"
      },
      "source": [
        "## Optional: Optimizer `config`s\n",
        "\n",
        "The `tensorflow_models` package defines serializable `config` classes that describe how to build the live objects. Earlier in this tutorial, you built the optimizer manually.\n",
        "\n",
        "The configuration below describes an (almost) identical optimizer built by the `optimizer_factory.OptimizerFactory`:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Fdb9C1ontnH_"
      },
      "outputs": [],
      "source": [
        "optimization_config = tfm.optimization.OptimizationConfig(\n",
        "    optimizer=tfm.optimization.OptimizerConfig(\n",
        "        type = \"adam\"),\n",
        "    learning_rate = tfm.optimization.LrConfig(\n",
        "        type='polynomial',\n",
        "        polynomial=tfm.optimization.PolynomialLrConfig(\n",
        "            initial_learning_rate=2e-5,\n",
        "            end_learning_rate=0.0,\n",
        "            decay_steps=num_train_steps)),\n",
        "    warmup = tfm.optimization.WarmupConfig(\n",
        "        type='linear',\n",
        "        linear=tfm.optimization.LinearWarmupConfig(warmup_steps=warmup_steps)\n",
        "    ))\n",
        "\n",
        "\n",
        "fac = tfm.optimization.optimizer_factory.OptimizerFactory(optimization_config)\n",
        "lr = fac.build_learning_rate()\n",
        "optimizer = fac.build_optimizer(lr=lr)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "Rp7R1hBfv5HG"
      },
      "outputs": [],
      "source": [
        "x = tf.linspace(0, num_train_steps, 1001).numpy()\n",
        "y = [lr(xi) for xi in x]\n",
        "plt.plot(x,y)\n",
        "plt.xlabel('Train step')\n",
        "plt.ylabel('Learning rate')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ywn5miD_dnuh"
      },
      "source": [
        "The advantage to using `config` objects is that they don't contain any complicated TensorFlow objects, and can be easily serialized to JSON, and rebuilt. Here's the JSON for the above `tfm.optimization.OptimizationConfig`:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "zo5RV5lud81Y"
      },
      "outputs": [],
      "source": [
        "optimization_config = optimization_config.as_dict()\n",
        "optimization_config"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "Z6qPXPEhekkd"
      },
      "source": [
        "The `tfm.optimization.optimizer_factory.OptimizerFactory` can just as easily build the optimizer from the JSON dictionary:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "p-bYrvfMYsxp"
      },
      "outputs": [],
      "source": [
        "fac = tfm.optimization.optimizer_factory.OptimizerFactory(\n",
        "    tfm.optimization.OptimizationConfig(optimization_config))\n",
        "lr = fac.build_learning_rate()\n",
        "optimizer = fac.build_optimizer(lr=lr)"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "name": "fine_tune_bert.ipynb",
      "private_outputs": true,
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
